import torch.nn as nn


class HeatmapDecoder(nn.Module):
    """ Heatmap decoder. """

    def __init__(self):
        super(HeatmapDecoder, self).__init__()
        self.decoder = nn.Sequential(
            nn.Conv2d(32, 16, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 8, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(8),
            nn.ReLU(inplace=True),
            nn.Conv2d(8, 2, kernel_size=1, stride=1, padding=0),
            nn.ReLU(inplace=True),
        )

    def forward(self, input_features):
        # Iterate til output block
        out = self.decoder(input_features)

        return out